"""miminal script that trains an OSEM varnet on simulated brainweb data
"""

from __future__ import annotations

import argparse
import json
from datetime import datetime
import utils
import parallelproj
import torch

from layers import EMUpdateModule
from models import Unet3D, SimpleOSEMVarNet, PostReconNet
from data import load_brain_image_batch, simulate_data_batch, download_brainweb_data

from pathlib import Path

parser = argparse.ArgumentParser(description='OSEM-VARNet reconstruction')
parser.add_argument('--num_datasets', type=int, default=12)
parser.add_argument('--num_training', type=int, default=40)
parser.add_argument('--num_validation', type=int, default=20)
parser.add_argument('--num_subsets', type=int, default=4)
parser.add_argument('--depth', type=int, default=8)
parser.add_argument('--num_epochs', type=int, default=500)
parser.add_argument('--num_epochs_post', type=int, default=500)
parser.add_argument('--batch_size', type=int, default=10)
parser.add_argument('--num_features', type=int, default=32)
parser.add_argument('--num_rings', type=int, default=4)
parser.add_argument('--radial_trim', type=int, default=181)
parser.add_argument('--random_seed', type=int, default=1)
parser.add_argument('--sens', type=float, default=1)
parser.add_argument('--voxel_size',
                    nargs='+',
                    type=float,
                    default=[2.5, 2.5, 2.66])
parser.add_argument('--fusion_mode', type=str, default = 'simple', choices=['simple', 'de_pierro'])

args = parser.parse_args()

num_datasets = args.num_datasets
num_training = args.num_training
num_validation = args.num_validation
num_subsets = args.num_subsets
depth = args.depth
num_epochs = args.num_epochs
num_epochs_post = args.num_epochs_post
batch_size = args.batch_size
num_features = args.num_features
num_rings = args.num_rings
radial_trim = args.radial_trim
random_seed = args.random_seed
sens = args.sens
voxel_size = tuple(args.voxel_size)
fusion_mode = args.fusion_mode

# device variable (cpu or cuda) that determines whether calculations
# are performed on the cpu or cuda gpu
if parallelproj.cuda_present:
    dev = 'cuda'
else:
    dev = 'cpu'

output_dir = Path(
    'run_osem_varnet') / f'{datetime.now().strftime("%Y%m%d_%H%M%S")}'
output_dir.mkdir(exist_ok=True, parents=True)

with open(output_dir / 'input_cfg.json', 'w') as f:
    json.dump(vars(args), f)

#----------------------------------------------------------------------------
#----------------------------------------------------------------------------
#--- setup the scanner / LOR geometry ---------------------------------------
#----------------------------------------------------------------------------
#----------------------------------------------------------------------------

# setup a line of response descriptor that describes the LOR start / endpoints of
# a "narrow" clinical PET scanner with 9 rings
lor_descriptor = utils.DemoPETScannerLORDescriptor(torch,
                                                   dev,
                                                   num_rings=num_rings,
                                                   radial_trim=radial_trim)
axial_fov_mm = float(lor_descriptor.scanner.num_rings *
                     (lor_descriptor.scanner.ring_positions[1] -
                      lor_descriptor.scanner.ring_positions[0]))

#----------------------------------------------------------------------------
#----------------------------------------------------------------------------
#--- load the brainweb images -----------------------------------------------
#----------------------------------------------------------------------------
#----------------------------------------------------------------------------

# download and extract the brainweb PET/MR images into ./data if not present
download_brainweb_data()

# image properties
ids = tuple([i for i in range(num_datasets)])

emission_image_database, attenuation_image_database = load_brain_image_batch(
    ids,
    torch,
    dev,
    voxel_size=voxel_size,
    axial_fov_mm=0.95 * axial_fov_mm,
    verbose=False)

img_shape = tuple(emission_image_database.shape[2:])

#----------------------------------------------------------------------------
#----------------------------------------------------------------------------

subset_projectors = parallelproj.SubsetOperator([
    utils.RegularPolygonPETProjector(
        lor_descriptor,
        img_shape,
        voxel_size,
        views=torch.arange(i,
                           lor_descriptor.num_views,
                           num_subsets,
                           device=dev)) for i in range(num_subsets)
])

#--------------------------------------------------------------------------------
#--------------------------------------------------------------------------------
#--------------------------------------------------------------------------------
#--------------------------------------------------------------------------------

print(f'simulating emission and correction data')

# simulate all emission and correction sinograms we need
emission_data_database, correction_database, contamination_database, adjoint_ones_database = simulate_data_batch(
    emission_image_database,
    attenuation_image_database,
    subset_projectors,
    sens=sens,
    random_seed=random_seed, xp=torch)

#--------------------------------------------------------------------------------
#--------------------------------------------------------------------------------
#--------------------------------------------------------------------------------
#--------------------------------------------------------------------------------

# run OSEM reconstructions of the simulated data

osem_update_modules = [
    EMUpdateModule(projector) for projector in subset_projectors.operators
]

osem_database = torch.ones((num_datasets, 1) + subset_projectors.in_shape,
                           device=dev,
                           dtype=torch.float32)

num_osem_iter = 102 // num_subsets

subset_order = utils.distributed_subset_order(num_subsets)

for i in range(num_osem_iter):
    print(f'OSEM iteration {(i+1):003}/{num_osem_iter:003}', end='\r')
    for j in range(num_subsets):
        subset = subset_order[j]
        osem_database = osem_update_modules[subset](
            osem_database, emission_data_database[subset, ...],
            correction_database[subset, ...],
            contamination_database[subset, ...], adjoint_ones_database[subset,
                                                                       ...])

#--------------------------------------------------------------------------------
#--------------------------------------------------------------------------------
# model training
#--------------------------------------------------------------------------------
#--------------------------------------------------------------------------------

print('\npostrecon unet training\n')

post_recon_unet = PostReconNet(Unet3D(num_features=num_features).to(dev))
post_recon_unet.train()

loss_fn_post = torch.nn.MSELoss()
optimizer_post = torch.optim.Adam(post_recon_unet.parameters(), lr=1e-3)

loss_arr_post = torch.zeros(num_epochs_post)

min_val_loss_post = float('inf')

for epoch in range(num_epochs_post):
    batch_inds = torch.split(torch.randperm(num_training), batch_size)

    for ib, batch_ind in enumerate(batch_inds):
        x_fwd_post = post_recon_unet(osem_database[batch_ind, ...])
        loss_post = loss_fn_post(x_fwd_post, emission_image_database[batch_ind,
                                                                     ...])

        print(
            f'{(epoch+1):03}/{num_epochs_post:03} {(ib+1):03} {loss_post:.2E}',
            end='\r')

        # Backpropagation
        loss_post.backward()
        optimizer_post.step()
        optimizer_post.zero_grad()

    loss_arr_post[epoch] = loss_post

    if (epoch + 1) % 20 == 0:
        post_recon_unet.eval()
        x_fwd_post = post_recon_unet(
            osem_database[num_training:(num_training + num_validation), ...])
        val_loss_post = loss_fn_post(
            x_fwd_post, emission_image_database[num_training:(num_training +
                                                              num_validation),
                                                ...])
        print(
            f'{(epoch+1):03}/{num_epochs_post:03} train_loss {float(loss_post):.2E} val_loss {val_loss_post:.2E}'
        )
        post_recon_unet.train()

        if val_loss_post < min_val_loss_post:
            min_val_loss_post = val_loss_post
            torch.save(post_recon_unet.neural_net.state_dict(),
                       output_dir / 'post_recon_model_best_state.pt')

torch.save(loss_arr_post, output_dir / 'training_loss_post.pt')
torch.save(post_recon_unet.neural_net.state_dict(),
           output_dir / 'post_recon_model_last_state.pt')

del post_recon_unet

#--------------------------------------------------------------------------------
#--------------------------------------------------------------------------------
# model training
#--------------------------------------------------------------------------------
#--------------------------------------------------------------------------------

print('\nvarnet training\n')

unet = Unet3D(num_features=num_features).to(dev)

if fusion_mode == 'simple':
    best_path = output_dir / 'post_recon_model_best_state.pt'
    print(f'loading pre-trained weights from {best_path}')
    unet.load_state_dict(torch.load(best_path))

osem_var_net = SimpleOSEMVarNet(osem_update_modules, unet, depth, dev, fusion_mode=fusion_mode)
print(f'fusion mode: {osem_var_net.fusion_mode}')
osem_var_net.train()

loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.Adam(osem_var_net.parameters(), lr=1e-3)

loss_arr = torch.zeros(num_epochs)
min_val_loss = float('inf')

for epoch in range(num_epochs):
    batch_inds = torch.split(torch.randperm(num_training), batch_size)

    for ib, batch_ind in enumerate(batch_inds):
        x_fwd = osem_var_net(osem_database[batch_ind, ...],
                             emission_data_database[:, batch_ind, ...],
                             correction_database[:, batch_ind, ...],
                             contamination_database[:, batch_ind, ...],
                             adjoint_ones_database[:, batch_ind, ...])

        loss = loss_fn(x_fwd, emission_image_database[batch_ind, ...])

        print(f'{(epoch+1):03}/{num_epochs:03} {(ib+1):03} {loss:.2E}',
              end='\r')

        # Backpropagation
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

    loss_arr[epoch] = loss

    if (epoch + 1) % 20 == 0:
        osem_var_net.eval()
        val_loss = 0
        for iv in range(num_training, (num_training + num_validation)):
            x_fwd = osem_var_net(osem_database[iv:(iv + 1), ...],
                                 emission_data_database[:, iv:(iv + 1), ...],
                                 correction_database[:, iv:(iv + 1), ...],
                                 contamination_database[:, iv:(iv + 1), ...],
                                 adjoint_ones_database[:, iv:(iv + 1), ...])
            val_loss += float(
                loss_fn(x_fwd, emission_image_database[iv:(iv + 1), ...]))

        val_loss /= num_validation
        print(f'{(epoch+1):03}/{num_epochs:03} train_loss {float(loss):.2E} val_loss {val_loss:.2E} net_weight {float(osem_var_net.neural_net_weight):.2E}')
        osem_var_net.train()

        if val_loss < min_val_loss:
            min_val_loss = val_loss
            torch.save(osem_var_net.state_dict(),
                       output_dir / 'model_best_state.pt')

torch.save(loss_arr, output_dir / 'training_loss.pt')
torch.save(osem_var_net.state_dict(), output_dir / 'model_last_state.pt')
